import sys
sys.path.append(".")

import websocket
import json
import time

from device_protocol_sdk.abstract_device import AbstractDevice
from device_protocol_sdk.abstract_device import ActionItem
from device_protocol_sdk.model.device_status import DeviceStatus,MessageLevel
from device_protocol_sdk.pusher import DevicePusher
import asyncio

from typing import List, Any, Dict
import logging, math, json
import threading
import time
import numpy as np

import roslibpy
# import msgs.sportModeState.SportModeState as SportModeState
from msgs import sportModeState,lowState,unitree_type_define



logger = logging.getLogger("Go2")



class Go2Device(AbstractDevice):

    def __init__(self, go2device_config: Dict):
        self.ws_url = go2device_config.get("ws_url","localhost")
        self.ws_url_port = go2device_config.get("port",9090)
        self.client:roslibpy.Ros = None
        self.sportmodestate_subscriber:roslibpy.Topic = None
        self.lowstate_subscriber:roslibpy.Topic = None
        # 高层运动控制 publisher
        self.sport_publisher:roslibpy.Topic = None

        self.sportMosdeState:sportModeState.SportModeState = None
        self.lowState:lowState.LowState = None

        self.connected = False
        self.should_run = True
        

        self.device_id = None
        self.connection_str = None

        self.device_statuses = {}
        self.unitree_request=unitree_type_define.UnitreeRequest()

    @property
    def protocol_name(self) -> str:
        return "ros"
    
    def _get_device_key(self, device_id: int, connection_str: str) -> str:
        """生成设备的唯一标识key"""
        return f"{device_id}:{connection_str}"

    def _get_device_status(self, device_key: str) -> Dict:
        """获取指定设备的状态字典"""
        if device_key not in self.device_statuses:
            self.device_statuses[device_key] = {}
        return self.device_statuses[device_key]

    def _update_device_status(self, device_key: str, updates: Dict):
        """更新指定设备的状态"""
        if device_key not in self.device_statuses:
            self.device_statuses[device_key] = {}
        if not updates:
            self.device_statuses[device_key] = {}
        else:
            self.device_statuses[device_key].update(updates)

    def _status_loop(self,connection_client, device_id: int, connection_str: str):
        """状态获取循环线程"""
        device_key = self._get_device_key(device_id, connection_str)
        logger.info(f"Starting status thread for device: {device_key}")
        # 初始化设备状态
        self._update_device_status(device_key, {})
        while not self._stop_event.is_set():
            self._get_go2_status(connection_client,device_id, connection_str)
            time.sleep(1)

        logger.info(f"Status thread stopped for device: {device_key}")

        if self.client.is_connected:
            self.client.terminate()
            self.sportmodestate_subscriber.unsubscribe()

    def _create_client(self, device_id: int, connection_str: str) -> tuple[bool,Any]:
        print("Creating Go2 ROS Client")
        self.device_id = device_id
        self.connection_str = connection_str
        if self.connection_str:
            self.ws_url = self.connection_str
        self.client = roslibpy.Ros(host=self.ws_url,port=self.ws_url_port)
        self.client.run()
        self.sportmodestate_subscriber = roslibpy.Topic(self.client, 'lf/sportmodestate', 'unitree_go/msg/SportModeState')
        self.lowstate_subscriber = roslibpy.Topic(self.client,'/lowstate','unitree_go/msg/LowState')
        self.sport_publisher = roslibpy.Topic(self.client, '/api/sport/request', 'unitree_api/msg/Request')
        self.sportmodestate_subscriber.subscribe(callback=self.sport_mode_state_sub_callback)
        self.lowstate_subscriber.subscribe(callback=self.low_state_sub_callback)

        #* Status loop thread
        self._stop_event = threading.Event()
        self._status_thread = threading.Thread(target=self._status_loop, args=(self.client,device_id, connection_str))
        self._status_thread.start()

        return True,self.client

    def _close_client(self, client, device_id: int, connection_str: str) -> Any:
        self.sportmodestate_subscriber.unsubscribe()
        client.terminate()
        return True
    
    def _get_go2_status(self,client:roslibpy.Ros, device_id: int, connection_str: str):
        device_key = self._get_device_key(device_id, connection_str)
        if not client:
            return 
        if not client.is_connected:
            logger.warning("ROS Bridge is not connected")
            return
        if self.sportMosdeState is None:
            return
        
        status_update = {}
        status_update['mode'] = self.sportMosdeState.mode
        status_update['gait_type'] = self.sportMosdeState.gait_type
        status_update['foot_raise_height'] = self.sportMosdeState.foot_raise_height
        status_update['position'] = self.sportMosdeState.position
        status_update['body_height'] = self.sportMosdeState.body_height
        status_update['velocity'] = self.sportMosdeState.velocity
        status_update['yaw_speed'] = self.sportMosdeState.yaw_speed
        status_update['range_obstacle'] = self.sportMosdeState.range_obstacle

        self._update_device_status(device_key, status_update)
        self._update_device_status(device_key, {'last_update': time.time()})
        

    def sport_mode_state_sub_callback(self,message):
        self.sportMosdeState = sportModeState.SportModeState(message)
    def low_state_sub_callback(self,message):
        self.lowState = lowState.LowState(message)

    def get_device_status(self, client, device_id: str, connection_str: str) -> DeviceStatus:
        """获取设备状态"""

        device_key = self._get_device_key(int(device_id), connection_str)
        status_data = self._get_device_status(device_key)
        # print("Status Data:", status_data)

        last_update = status_data.get('last_update',0)
        if time.time() - last_update > 5:
            logger.warning(f"Device {device_key} status is stale")
            
            device_status = DeviceStatus(is_lock=True,
                                         heartbeat=False,alt=0,
                                         battery=0,lat=0,lon=0,
                                         airspeed=0,
                                         groundspeed=0,
                                         roll=0,pitch=0,yaw=0,height=0)
            return device_status
        else:
            return DeviceStatus(
            is_lock=False,
            heartbeat=True,
            alt=0,
            battery=self.lowState.battery_state['soc'],
            lat=0,
            lon=0,
            airspeed=0,
            groundspeed=0,
            roll=0,
            pitch=0,
            yaw=0,
            height=self.sportMosdeState.body_height
            )

        device_key = self._get_device_key(int(device_id), connection_str)
        status_data = self._get_device_status(device_key)
        print("Status Data:", status_data)
        if not status_data:
            return None
        # 检查状态是否过期（超过5秒没有更新）
        last_update = status_data.get('last_update', 0)
        if time.time() - last_update > 5:
            logger.warning(f"Device {device_key} status is stale")
        return DeviceStatus(
            is_lock=status_data.get('is_lock', 1) == 0,  # 转换为布尔值
            heartbeat=status_data.get('heartbeat', 1) == 0,
            alt=status_data.get('alt', 0),
            battery=status_data.get('battery', 0),
            lat=status_data.get('lat', 0),
            lon=status_data.get('lon', 0),
            # 添加更多状态字段
            airspeed=status_data.get('airspeed', 0),
            groundspeed=status_data.get('groundspeed', 0),
            roll=status_data.get('roll', 0),
            pitch=status_data.get('pitch', 0),
            yaw=status_data.get('yaw', 0),
            height=status_data.get('height', 0)
        )


    def execute(self, client, device_id: int, connection_str: str, command_type: str, params: Dict[str, Any]) -> Dict[
        str, Any]:
        if command_type == 'standdown':
            return self.standdown()
        elif command_type == 'standup':
            return self.standup()
        else:
            raise NotImplementedError(f"不支持的指令类型: {command_type}")

    def get_action_list(self) -> List[ActionItem]:
        return [
            ActionItem(
                name="站立",
                command_type="standup",
                description="机器狗站立起来,关节锁定，站高。",
                params={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            ActionItem(
                name="趴下",
                command_type="standdown",
                description="机器狗趴下，电机关节保持锁定",
                params={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            ActionItem(
                name="解除锁定",
                command_type="balancestand",
                description="解除关节电机锁定，从正常站立模式切换到平衡站立模式。"
                            "这一模式下机身姿态和高度会始终保持平衡，不受地形的影响。"
                            "可以通过调用 Euler()和 BodyHeight()接口实现机身的字体和高度控制(详见表格的对应部分)。",
                params={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            ActionItem(
                name="进入阻尼状态",
                command_type="damp",
                description="所有电机关节停止运动并进入阻尼状态。该模式具有最高的优先级，用于突发情况下的急停。",
                params={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            ActionItem(
                name="恢复站立",
                command_type="recoverystand",
                description="从翻倒或趴下状态恢复至平衡站立状态。不论是否翻倒，都会恢复至站立。",
                params={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            ActionItem(
                name="站立和行走时的姿态",
                command_type="euler",
                description="设置 Go2 平衡站立或移动时的机体姿态角。欧拉角采用绕机体相对轴和 z-y-x 旋转顺序的表示方式。",
                params={
                    "type": "object",
                    "properties": {
                        "roll": {
                            "type": "number",
                            "description": "横滚角，单位为度，取值范围:[-0.75~0.75] (rad)"
                        },
                        "pitch": {
                            "type": "number",
                            "description": "俯仰角，单位为度，取值范围:[-0.75~0.75] (rad)"
                        },
                        "yaw": {
                            "type": "number",
                            "description": "偏航角，单位为度，取值范围:[-0.6~0.6] (rad)"
                        }
                    },
                    "required": ["roll","pitch","yaw"]
                }
            ),
            ActionItem(
                name="移动",
                command_type="move",
                description="控制移动速度。设定的速度为机体坐标系表示下的速度。",
                params={
                    "type": "object",
                    "properties": {
                        "vx": {
                            "type": "number",
                            "description": "机体坐标系下x轴方向的速度，取值范围[-2.5~3.8] (m/s)"
                        },
                        "vy": {
                            "type": "number",
                            "description": "机体坐标系下y轴方向的速度，取值范围[-1.0~1.0] (m/s)"
                        },
                        "vyaw": {
                            "type": "number",
                            "description": "机体坐标系下绕yaw轴方向旋转的速度，取值范围[-4~4] (rad/s)"
                        }
                    },
                    "required": ["vx","vy","vyaw"]
                }
            ),
            ActionItem(
                name="坐下",
                command_type="sit",
                description="特殊动作，机器狗坐下。需要注意特殊动作需要在上一个动作执行完毕后再执行，否则会导致动作异常。",
                params={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            ActionItem(
                name="站起（相对于坐下）",
                command_type="risesit",
                description="从坐下状态恢复到平衡站立",
                params={
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            ),
            ActionItem(
                name="切换步态",
                command_type="switchgait",
                description="切换机器人的步态以应对不同的路况或用户喜好。其中：正向爬楼模式为机器人头朝楼梯，逆向爬楼模式为机器人背朝楼梯",
                params={
                    "type": "object",
                    "properties": {
                        "d": {
                            "type": "number",
                            "description": "步态枚举值，取值：0~4，其中0为漫步， 1为小跑，2为快跑 ，3为正向爬楼模式，4为逆向爬楼模式。"
                        }
                    },
                    "required": ["d"]
                }
            )
        ]


    def standdown(self):
        print("standdown")
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_STANDDOWN)
        return {'status':'success'}
    def standup(self):
        print("standup")
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_STANDUP)
        return {'status':'success'}

    def balancestand(self):
        print("BalanceStand")
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_BALANCESTAND)
        return {'status':'success'}
    def damp(self):
        print("damp")
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_AI_ID_DAMP)
        return {'status':'success'}

    def recoverystand(self):
        print("recoveryStand")
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_RECOVERYSTAND)
        return {'status':'success'}

    def euler(self,roll:float=0,pitch:float=0,yaw:float=0):
        print("euler")
        parameter = json.dumps({'roll':roll,'pitch':pitch,'yaw':yaw})
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_EULER,parameter=parameter)
        return {'status':'success'}
    
    def move(self,vx:float=0,vy:float=0,vyaw:float=0):
        print("move")
        parameter = json.dumps({'vx':vx,'vy':vy,'vyaw':vyaw})
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_MOVE,parameter=parameter)
        return {'status':'success'}
    
    def sit(self):
        print("sit")
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_SIT)
        return {'status':'success'}
    def risesit(self):
        print("riseSit")
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_RISESIT)
        return {'status':'success'}
    def switchgait(self,d:int=0):
        print("switchGait")
        parameter = json.dumps({'d':d})
        self.sport_api(api_id=unitree_type_define.ROBOT_SPORT_API_ID_SWITCHGAIT,parameter=parameter)
        return {'status':'success'}


    def sport_api(self,api_id:int=None,parameter:str=None,binary:list=None):
        if api_id is not None:
            self.unitree_request.setRequestIdentity(api_id=api_id)
        self.unitree_request.setRequestHeader()
        req_ = self.unitree_request.Request(parameter=parameter,binary=binary,if_header_set=True)
        print("Sport API Request:",req_)
        msg = roslibpy.Message(req_)
        self.sport_publisher.publish(msg)


async def main():
    device_config = {
        "ws_url": "localhost",
        "port": 9090
    }
    pusher = DevicePusher(lambda: Go2Device(device_config))
    await pusher.connect_server(
        grpc_server_url="192.168.50.167:50058",
        type_description="机器狗go2"
    )

if __name__ == "__main__":
    
    asyncio.run(main())
    
    
    # go2_device = Go2Device(device_config)
    #
    # # go2_device.ws.run_forever()
    #
    # reconnect_interval = 5  # 断线后等待5秒重试
    # while True:
    #     try:
    #         # 启动WebSocket客户端
    #         go2_device.ws.run_forever()
    #     except Exception as e:
    #         print(f"发生异常: {e}")
    #         print(f"{reconnect_interval}秒后尝试重连...")
    #         time.sleep(reconnect_interval)